import matplotlib.pyplot as plt
from collections import Counter
from io import BytesIO
import pandas as pd
import seaborn as sns
import json
import os

class MOOCSentimentVisualizer:
    def __init__(self, data):
        self.set_chinese_font()
        self.df = pd.DataFrame(data)

    def set_chinese_font(self):
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

    def plot_final_sentiment_distribution(self, save_path=None):
        sentiment_counts = Counter(self.df['最终情感'])
        fig, ax = plt.subplots(figsize=(8, 6))
        key_mapping = {
            'positive': '积极面',
            'mixed': '中立',
            'negative': '消极面'
        }
        # 创建新的字典，通过替换键名
        sentiment_counts = Counter({key_mapping.get(k, k): v for k, v in sentiment_counts.items()})
        sns.barplot(x=list(sentiment_counts.keys()), y=list(sentiment_counts.values()), ax=ax, palette="viridis", hue=list(sentiment_counts.keys()), legend=False)
        ax.set_title("MOOC 课程情感分布")
        ax.set_xlabel("情感类别")
        ax.set_ylabel("频率")
        return self.save_or_return_image(fig, save_path)

    def plot_aspect_frequency(self, save_path=None):
        aspect_terms = [term.strip() for aspect_list in self.df['方面词'] for term in aspect_list.split(',')]
        aspect_counts = Counter(aspect_terms)
        aspect_counts = {word: count for word, count in aspect_counts.items() if count >= 2}
        fig, ax = plt.subplots(figsize=(10, 8))
        sns.barplot(x=list(aspect_counts.keys()), y=list(aspect_counts.values()), ax=ax, palette="coolwarm", hue=list(aspect_counts.keys()), legend=False)
        ax.set_title("用户评价的方面词频率")
        ax.set_xlabel("方面词")
        ax.set_ylabel("频率")
        # plt.xticks(rotation=45, ha='right')  # 旋转标签
        return self.save_or_return_image(fig, save_path)

    def plot_sentiment_word_frequency(self, save_path=None):
        sentiment_words = [word.strip() for word_list in self.df['情感词'] for word in word_list.split(',')]
        sentiment_word_counts = Counter(sentiment_words)
        sentiment_word_counts = {word: count for word, count in sentiment_word_counts.items() if count > 2}
        fig, ax = plt.subplots(figsize=(15, 8))
        sns.barplot(x=list(sentiment_word_counts.keys()), y=list(sentiment_word_counts.values()), ax=ax, palette="magma", hue=list(sentiment_word_counts.keys()), legend=False)
        ax.set_title("情感词频率分布")
        ax.set_xlabel("情感词")
        ax.set_ylabel("频率")
        # plt.xticks(rotation=45, ha='right')  # 旋转标签
        return self.save_or_return_image(fig, save_path)

    def save_or_return_image(self, fig, save_path):
        # 如果指定了保存路径，则保存图像
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            fig.savefig(save_path, format="png")
            plt.close(fig)
            print(f"图像已保存到: {save_path}")
        else:
            # 否则将图像返回为 BytesIO 对象
            buf = BytesIO()
            fig.savefig(buf, format="png")
            buf.seek(0)
            plt.close(fig)
            return buf
